import io
import json
import os
from flask import Flask, request, jsonify, Response, render_template
import base64
import cv2
from io import BytesIO
from PIL import Image
import numpy as np
import dlib
import face_recognition
from flask_cors import CORS
from datetime import datetime
from sqlalchemy import create_engine, Table, Column, Integer, DateTime, String, Enum, MetaData, LargeBinary, insert
from sqlalchemy.orm import scoped_session, sessionmaker
import uuid
app = Flask(__name__, static_folder='static', template_folder='templates')
CORS(app)
app.config['TEMPLATES_AUTO_RELOAD'] = True
app.config['SEND_FILE_MAX_AGE_DEFAULT'] = 0
app.config['SQLALCHEMY_DATABASE_URI'] = 'mysql://root@localhost:3306/flask_sql'

# 创建 MySQL 数据库连接
engine = create_engine('mysql+pymysql://root:root@localhost:3306/flask_sql', echo=True)
# 定义数据表
metadata = MetaData()
dan = Table('dan', metadata,
            Column('id', Integer, primary_key=True),
            Column('time', DateTime, nullable=False),
            Column('location', String(64), nullable=False),
            Column('behavior', String(64), nullable=False),
            Column('image_data', LargeBinary, nullable=False),
            Column('status', Enum('pending', 'approved', 'rejected'), nullable=False)
            )

# 创建数据表
metadata.create_all(engine)

Session = scoped_session(sessionmaker(bind=engine))
detector = dlib.get_frontal_face_detector()
sp = dlib.shape_predictor('./dat/shape_predictor_68_face_landmarks.dat')
facerec =dlib.face_recognition_model_v1('./dat/dlib_face_recognition_resnet_model_v1.dat')
registered_faces_file = 'registered_faces.json'

# 创建截图保存目录
#save_directory = "screenshots"
#os.makedirs(save_directory, exist_ok=True)

def load_registered_faces():
    if os.path.exists(registered_faces_file):
        with open(registered_faces_file, 'r') as file:
            return json.load(file)
    return {}

registered_faces = load_registered_faces()

# 全局变量：定义摄像头对象
camera = None
# 处理关闭摄像头的请求
def close_camera():
    global camera
    if camera is not None:
        camera.release()  # 释放摄像头资源
        camera = None
        return jsonify({'message': '摄像头已关闭'})
    else:
        return jsonify({'message': '摄像头未打开'})

#全局变量，记录陌生人识别次数
stranger_count = 0

@app.route('/video_feed/0')
def video_feed():
    try:
        stream_url = 0  # 0代表第一个本地摄像头

        def generate_frames():
            global camera
            global stranger_count

            if isinstance(stream_url, int):  # 处理本地摄像头
                camera = cv2.VideoCapture(stream_url)
            else:
                camera = cv2.VideoCapture(stream_url)

            # 设置分辨率为480*480
            camera.set(cv2.CAP_PROP_FRAME_WIDTH, 480)
            camera.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

            frame_skip = 5  # 跳过的帧数
            frame_count = 0

            while True:
                success, frame = camera.read()
                if not success:
                    break

                if frame_count % frame_skip == 0:
                    # 1. 转换为灰度图像
                    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
                    # 2. 检测人脸
                    faces = detector(gray, 1)
                    stranger_detected = False  # 标记是否检测到陌生人

                    for face in faces:
                        # 获取人脸关键点
                        shape = sp(frame, face)
                        # 计算人脸的128维编码
                        face_encoding = np.array(facerec.compute_face_descriptor(frame, shape))
                        # 比较捕获的人脸与已注册人脸库中的编码，以判断是否为已知人脸
                        matches = face_recognition.compare_faces(list(registered_faces.values()), face_encoding,
                                                                 tolerance=0.4)
                        name = "Stranger"
                        color = (0, 0, 255)  # 默认红色标记陌生人

                        if True in matches:
                            first_match_index = matches.index(True)
                            staff_name = list(registered_faces.keys())[first_match_index]
                            name = staff_name
                            color = (0, 255, 0)  # 绿色标记已注册人脸
                        else:
                            stranger_detected = True

                        # 在人脸周围绘制矩形框
                        cv2.rectangle(frame, (face.left(), face.top()), (face.right(), face.bottom()), color, 2)

                        # 添加文本标签
                        cv2.putText(frame, name, (face.left(), face.top() - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color,
                                    2)

                    # 3. 编码图像
                    ret, buffer = cv2.imencode('.jpg', frame, [int(cv2.IMWRITE_JPEG_QUALITY), 70])
                    frame_bytes = buffer.tobytes()
                    yield (b'--frame\r\n'b'Content-Type: image/jpeg\r\n\r\n' + frame_bytes + b'\r\n')  # 构造响应体

                    # 4. 如果是陌生人，保存图片并插入数据库
                    if stranger_detected:
                        stranger_count += 1  # 增加陌生人计数

                        if isinstance(frame, np.ndarray):
                            _, image_binary = cv2.imencode('.jpg', frame)
                            image_data = image_binary.tobytes()

                            # 获取当前时间
                            current_time = datetime.now()

                            # 使用 SQLAlchemy 插入数据
                            session = Session()
                            try:
                                stmt = dan.insert().values(
                                    time=current_time,
                                    location='系统',
                                    behavior='陌生人员',
                                    image_data=image_data,
                                    status='pending'
                                )
                                session.execute(stmt)
                                session.commit()
                                print("数据插入成功")

                            except Exception as e:
                                session.rollback()
                                print(f"插入数据失败: {e}")
                            finally:
                                session.close()
                        else:
                            print(f"Frame is not a valid numpy array: {type(frame)}")

                frame_count += 1

        return Response(generate_frames(), mimetype='multipart/x-mixed-replace; boundary=frame')

    except Exception as e:
        app.logger.error(f"Failed to stream video for camera: {str(e)}")
        return jsonify({'error': 'Failed to process video stream'}), 500


@app.route('/close_camera', methods=['POST'])
def handle_close_camera():
    data = request.get_json()
    action = data.get('action')

    app.logger.info(f"Received action: {action}")

    if action == 'close_camera':
        response = close_camera()
        app.logger.info("Camera close response: %s", response.json)
        return response
    else:
        app.logger.error("Unknown action received")
        return jsonify({'error': '未知操作'}), 400

@app.route('/')
def index():
    return render_template('index.html')

if __name__ == '__main__':
    app.run(debug=True)